package com;

import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;

import java.time.Duration;
import java.util.*;
import java.util.stream.Collectors;

import static com.BasicProducer.BROKERS;
import static com.BasicProducer.TOPIC_NAME;

/**
 * @author 罗俊华
 * @date 2021/12/16 - 4:17 下午
 */
@Slf4j
public class BasicConsumer {


    public static final String CONSUMER_GROUP = "group-java";

    @DisplayName("自动提交")
    @Test
    public void autoCommit() {

        KafkaConsumer<String, String> consumer = prepareConsumer(true);
//        1、消费者订阅的主题列表
        consumer.subscribe(Collections.singletonList(TOPIC_NAME));


        for (; ; ) {

            ConsumerRecords<String, String> poll = consumer.poll(Duration.ofSeconds(10));

            log.info("收到{}条批量消息", poll.count());

            for (ConsumerRecord<String, String> record : poll) {

                log.info("主题：{}，分区：{}，偏移量：{}消息内容：{}", record.topic(), record.partition(), record.offset(), record.value());
            }


        }

    }

    @Test
    @DisplayName("手动提交")
    public void manualCommit() {

        KafkaConsumer<String, String> consumer = prepareConsumer(false);
//        1、消费者订阅的主题列表
        consumer.subscribe(Collections.singletonList(TOPIC_NAME));

        for (; ; ) {

            //
            ConsumerRecords<String, String> poll = consumer.poll(Duration.ofSeconds(10));

            log.info("收到{}条批量消息", poll.count());

            for (ConsumerRecord<String, String> record : poll) {

                log.info("主题：{}，分区：{}，偏移量：{}消息内容：{}", record.topic(), record.partition(), record.offset(), record.value());
            }


            if (poll.count() > 0) {
                log.info("轮询收到的消息已经全部消费完");


                log.info("执行手动同步提交");
                // 在消息消费完之后调用同步提交的方法，集群返回ack之前就一直阻塞
//                  返回 ack 之后表示提交成功
                // 手动提交 offset，当前线程会被阻塞，直到 offset 提交成功
                consumer.commitSync();


//                手动异步提交：在消息消费完之后提交，不需要等待集群的ack
//                直接执行之后的逻辑，供集群调用
                /*log.info("执行消息异步手动提交");
                consumer.commitAsync((Map<TopicPartition, OffsetAndMetadata> offsets, Exception exception) -> {

                    if (exception != null) {
                        log.error("异步手动提交失败", exception);
                    }

                    offsets.entrySet().forEach((Map.Entry<TopicPartition, OffsetAndMetadata> entry) -> {

                        TopicPartition key = entry.getKey();
                        OffsetAndMetadata value = entry.getValue();
                        log.info("主题：{}，分区：{}，offset：{}，metadata：{}", key.topic(), key.partition(), value.offset(), value.metadata());
                    });
                });*/
            }

        }


    }

    @DisplayName("消费者指定分区进行消费")
    @Test
    public void consumeFromPartition() {
        KafkaConsumer<String, String> consumer = prepareConsumer(true);


//        2、消费者指定分区进行消费
        consumer.assign(Collections.singletonList(new TopicPartition(TOPIC_NAME, 1)));
//        // 从分区的起点，开始消费 （可选）
        consumer.seekToBeginning(Collections.singletonList(new TopicPartition(TOPIC_NAME, 1)));
        consumeMessage(consumer);

    }


    @DisplayName("指定 offset 消费")
    @Test
    public void consumeFromOffset() {


        KafkaConsumer<String, String> consumer = prepareConsumer(true);

//        3、指定 offset 消费
        consumer.assign(Collections.singletonList(new TopicPartition(TOPIC_NAME, 0)));
//        指定分区的offset，开始消费 （可选）
        consumer.seek(new TopicPartition(TOPIC_NAME, 0), 5);

        consumeMessage(consumer);

    }


    /**
     * 根据指定的时间去所有的 partition 中确定该时间对应的 offset
     * 然后去对应的 partition 中找到该 offset 及其之后的消息
     * 然后开始消费
     */
    @DisplayName("从指定的时间点开始消费")
    @Test
    public void consumeFromTime() {


        KafkaConsumer<String, String> consumer = prepareConsumer(true);

//          4、从指定的时间点开始消费
//        获取到该 topic 下的所有分区
        List<PartitionInfo> partitionInfos = consumer.partitionsFor(TOPIC_NAME);

//        从1小时前开始消费
        long fetchDateTime = new Date().getTime() - 1000 * 60 * 60;

        Map<TopicPartition, Long> topicPartitionMap = partitionInfos.stream().collect(Collectors.toMap((PartitionInfo partitionInfo) -> {
            return new TopicPartition(partitionInfo.topic(), partitionInfo.partition());
        }, (PartitionInfo partitionInfo) -> {
            return fetchDateTime;
        }));

//      查询每个分区从 fetchDateTime 时间之后的消息的 offset
        Map<TopicPartition, OffsetAndTimestamp> partitionOffsetAndTimestampMap = consumer.offsetsForTimes(topicPartitionMap);

        for (Map.Entry<TopicPartition, OffsetAndTimestamp> entry : partitionOffsetAndTimestampMap.entrySet()) {

            TopicPartition key = entry.getKey();

            OffsetAndTimestamp value = entry.getValue();

            if (key == null || value == null) {
                continue;
            }

            long offset = value.offset();

            log.info("从{}主题的{}分区的{}offset开始消费", key.topic(), key.partition(), value.offset());

            consumer.assign(Collections.singletonList(key));
            consumer.seek(key, offset);

        }

        consumeMessage(consumer);
    }

    /**
     * 执行消费消息
     *
     * @param consumer
     */
    private void consumeMessage(KafkaConsumer<String, String> consumer) {
        for (; ; ) {


            for (ConsumerRecord<String, String> consumerRecord : consumer.poll(Duration.ofMinutes(1))) {


                log.info("消费：主题：{}分区：{}偏移量：{}，消息的时间戳：{}，消息内容：{}", consumerRecord.topic(), consumerRecord.partition(), consumerRecord.offset(), new Date(consumerRecord.timestamp()), consumerRecord.value());

            }


        }
    }


    private KafkaConsumer<String, String> prepareConsumer(boolean autoCommit) {
        Properties properties = new Properties();

        properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, BROKERS);


        properties.put(ConsumerConfig.GROUP_ID_CONFIG, CONSUMER_GROUP);

//        是否自动提交 offset，默认就是 true
        properties.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, autoCommit ? "true" : "false");

//        自动提交 offset 的 时间间隔
        properties.put(ConsumerConfig.AUTO_COMMIT_INTERVAL_MS_CONFIG, "1000");

        /*
         * 当消费此主题的消费者是一个新的消费者组，
         * 或者指定 offset 的消费方式，offset 不存在
         * 那么应该如何消费
         * latest（默认）：只消费自己启动之后，发送到 topic 的消息
         * earliest：第一次从头开始消费，以后按照消费 offset 记录继续向后消费
         *             这个需要区别于 seekToBeginning （每次从头开始消费）
         * */
        properties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest");


        properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());


        // 一次 poll 最多拉取消息的条数，可以根据消费的速度来设置
        properties.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, 500);
//          如果两次 poll 的时间超过了 30s ，kafka 会认为其消费能力过弱
//        将其踢出消费者组。将分区分配给其他消费者
//        会执行 rebalance
        properties.put(ConsumerConfig.MAX_POLL_INTERVAL_MS_CONFIG, 30 * 1000);

//        consumer 给 broker 发送心跳的间隔时间
        properties.put(ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 1000);

//        kafka 如果超过 10s 都没有收到消费者的心跳，则会把消费者踢出消费者组
//        进入 rebalance，把分区分配给其他消费者
        properties.put(ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 10 * 1000);

        return new KafkaConsumer<>(properties);
    }

}
